{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain QA Application with Xinference and LangChain\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This demo walks through how to build an LLM-driven question-answering (QA) application with Xinference, Milvus, and LangChain."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy Xinference Locally or in a Distributed Cluster.\n",
    "\n",
    "For local deployment, run `xinference`. It will log an endpoint for you to use.\n",
    "\n",
    "To deploy Xinference in a cluster, first start an Xinference supervisor using the `xinference-supervisor`. You can also use the option -p to specify the port and -H to specify the host. The default port is 9997. If the default port is used, Xinference will choose an unused port for you. It will also log the endpoint for you to use.\n",
    "\n",
    "Then, start the Xinference workers using `xinference-worker` on each server you want to run them on. \n",
    "\n",
    "You can consult the README file from [Xinference](https://github.com/xorbitsai/inference) for more information.\n",
    "## Start a Model\n",
    "\n",
    "To use Xinference with LangChain, you need to first launch a model. You can use command line interface (CLI) to do so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uid: 19c73cee-3506-11ee-b286-fa163e74fa2d\n"
     ]
    }
   ],
   "source": [
    "!xinference launch --model-name \"falcon-instruct\" --model-format pytorch --size-in-billions 40 -e \"http://127.0.0.1:56256\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The command will return a model UID for you to use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "loader = TextLoader(\"/home/nijiayi/inference/examples/state_of_the_union.txt\") # Replace with the path of the document you want to query from\n",
    "\n",
    "documents = loader.load()\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = 512,\n",
    "    chunk_overlap  = 100,\n",
    "    length_function = len,\n",
    ")\n",
    "docs = text_splitter.split_documents(documents)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Up an Embedding Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings import XinferenceEmbeddings\n",
    "\n",
    "xinference_embeddings = XinferenceEmbeddings(\n",
    "    server_url=\"http://127.0.0.1:56256\", \n",
    "    model_uid = \"19c73cee-3506-11ee-b286-fa163e74fa2d\" # model_uid is the uid returned from launching the model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to the Vector Database"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For vector store, we use the Milvus vector database. [Milvus](https://milvus.io/docs/overview.md) is a database that stores, indexes, and manages massive embedding vectors generated by deep neural networks and other machine learning models. To run, you can first [Install Milvus Standalone with Docker Compose](https://milvus.io/docs/install_standalone-docker.md), or use Milvus Lite in the following way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "bat"
    }
   },
   "outputs": [],
   "source": [
    "$ pip install milvus\n",
    "\n",
    "$ milvus-server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Milvus\n",
    "\n",
    "vector_db = Milvus.from_documents(\n",
    "    docs,\n",
    "    xinference_embeddings,\n",
    "    connection_args={\"host\": \"0.0.0.0\", \"port\": \"19530\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query about the Document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
      "\n",
      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n"
     ]
    }
   ],
   "source": [
    "query = \"what does the president say about Ketanji Brown Jackson\"\n",
    "docs = vector_db.similarity_search(query, k=10)\n",
    "print(docs[0].page_content) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Inference Based on the Document"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we use Llama 2 Chat model supported by Xinference for inference task. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uid: 333e1d68-3507-11ee-a0d6-fa163e74fa2d\n"
     ]
    }
   ],
   "source": [
    "!xinference launch --model-name \"llama-2-chat\" --model-format ggmlv3 --size-in-billions 70 -e \"http://127.0.0.1:56256\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import Xinference\n",
    "\n",
    "xinference_llm = Xinference(\n",
    "    server_url=\"http://127.0.0.1:56256\",\n",
    "    model_uid = \"333e1d68-3507-11ee-a0d6-fa163e74fa2d\" # model_uid is the uid returned from launching the model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we can query the LLM without using the document:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nWhat did the president say about Ketanji Brown Jackson?\\nPresident Joe Biden called Judge Ketanji Brown Jackson a \"historic\" and \"inspiring\" nominee when he introduced her as his pick to replace retiring Supreme Court Justice Stephen Breyer. He highlighted her experience as a public defender and her commitment to justice and equality, saying that she would bring a unique perspective to the court.\\n\\nBiden also praised Jackson\\'s reputation for being a \"fair-minded\" and \"thoughtful\" jurist who is known for her ability to build'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xinference_llm(prompt=\"What did the president say about Ketanji Brown Jackson?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now query using the document to compare the result. We can create a memory object to track the chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create ConversationalRetrievalChain with chat model and the vectorstore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationalRetrievalChain\n",
    "\n",
    "chain = ConversationalRetrievalChain.from_llm(\n",
    "    llm=xinference_llm,\n",
    "    retriever=vector_db.as_retriever(),\n",
    "    memory=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can query information from the document. Instead of simply returning identical sentences from the document, the model generates responses by summarizing relevant content. Furthermore, it can relate a new query to the chat history, creating a chain of responses that build upon each other. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' According to the provided text, President Biden said that he nominated Circuit Court of Appeals Judge Ketanji Brown Jackson to serve on the United States Supreme Court 4 days ago, and that she is one of our nation’s top legal minds who will continue Justice Breyer’s legacy of excellence.'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "result = chain({\"question\": query})\n",
    "result[\"answer\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the LLM is capable of using the provided document to answer questions and summarize content. We can ask a few more questions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'  According to the given text, President Biden said that Ketanji Brown Jackson succeeded Justice Breyer on the Supreme Court.'"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"Did he mention who she succeeded\"\n",
    "result = chain({\"question\": query})\n",
    "result[\"answer\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The LLM accurately recognizes that \"he\" refers to \"the president\", and \"she\" refers to \"Ketanji Brown Jackson\" mentioned in the previous query. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'  According to the text, the president views COVID-19 as a \"God-awful disease\" and wants to move forward in addressing it in a unified manner, rather than allowing it to continue being a partisan dividing line.'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"Summarize the President's opinion on COVID-19\"\n",
    "result = chain({\"question\": query})\n",
    "result['answer']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the impressive capabilities of the LLM, and LangChain's \"chaining\" feature also allows for more coherent and context-aware interactions with the model."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
